#!/usr/bin/env python3.5
# sql中的变量值怎么处理？
from scapy.all import *
from scapy.layers.inet import TCP, IP
from datetime import datetime
import json
import sys
import re
import struct

def parse_parameters(payload):
    """解析参数化查询中的实际值（PostgreSQL 协议格式）"""
    try:
        # 示例：Bind 消息格式（'B'开头）
        if not payload.startswith('B'):
            return None
            
        # 跳过消息类型（1字节）和消息长度（4字节）
        pos = 5
        stmt_name = payload[pos:].split('\x00')[0]
        pos += len(stmt_name) + 1
        
        # 读取参数格式代码（2字节/参数）
        num_params = struct.unpack('!H', payload[pos:pos+2])[0]
        pos += 2
        param_formats = []
        for _ in range(num_params):
            param_formats.append(struct.unpack('!H', payload[pos:pos+2])[0])
            pos += 2
        
        # 读取参数值
        param_values = []
        for _ in range(num_params):
            val_len = struct.unpack('!I', payload[pos:pos+4])[0]
            pos += 4
            if val_len == 0xFFFFFFFF:  # NULL 值
                param_values.append(None)
            else:
                param_values.append(payload[pos:pos+val_len])
                pos += val_len
        
        return {
            'statement': stmt_name,
            'parameters': param_values
        }
    except Exception as e:
        print("参数解析错误: {}".format(e))
        return None

def parse_sql_with_params(payload):
    """组合预处理语句和实际参数"""
    try:
        # 解析 Parse 消息（'P'开头）
        if payload.startswith('P'):
            parts = payload[5:].split('\x00')
            stmt_name = parts[0]
            query = parts[1]
            return {
                'type': 'PREPARE',
                'name': stmt_name,
                'query': query
            }
        
        # 解析 Bind 消息（'B'开头）
        param_data = parse_parameters(payload)
        if param_data:
            return {
                'type': 'BIND',
                'data': param_data
            }
        
        # 简单查询直接提取
        sql_match = re.search(
            r'(SELECT|INSERT|UPDATE|DELETE|CREATE|ALTER|DROP)\s+.+?(?=\x00|;)', 
            payload, 
            re.IGNORECASE
        )
        return {'type': 'RAW_SQL', 'query': sql_match.group(0)} if sql_match else None
        
    except Exception as e:
        print("SQL解析错误: {}".format(e))
        return None

def extract_gaussdb_packets(pcap_file, db_port=5432):
    """主解析函数"""
    packets = rdpcap(pcap_file)
    results = []
    prepared_stmts = {}  # 缓存预处理语句 {stmt_name: query}
    
    for pkt in packets:
        try:
            if not pkt.haslayer(TCP) or pkt[TCP].dport != db_port:
                continue
                
            # 基础信息
            client_ip = pkt[IP].src if pkt.haslayer(IP) else "N/A"
            client_port = pkt[TCP].sport
            timestamp = datetime.fromtimestamp(float(pkt.time)).strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
            
            # 解析负载
            raw_payload = bytes(pkt[TCP].payload).decode('utf-8', errors='ignore')
            parsed = parse_sql_with_params(raw_payload)
            
            if not parsed:
                continue
                
            # 处理不同类型消息
            if parsed['type'] == 'PREPARE':
                prepared_stmts[parsed['name']] = parsed['query']
                continue
                
            elif parsed['type'] == 'BIND':
                stmt_name = parsed['data']['statement']
                if stmt_name in prepared_stmts:
                    query = prepared_stmts[stmt_name]
                    params = parsed['data']['parameters']
                    
                    # 将参数值插入到查询中（简化版）
                    final_sql = query
                    for i, val in enumerate(params, 1):
                        if val is None:
                            replacement = 'NULL'
                        else:
                            replacement = "'{}'".format(val.replace("'", "''"))
                        final_sql = final_sql.replace('${}'.format(i), replacement)
                    
                    results.append({
                        'timestamp': timestamp,
                        'client': "{}:{}".format(client_ip, client_port),
                        'type': 'PARAMETRIC',
                        'prepared_stmt': stmt_name,
                        'original_query': query,
                        'parameters': params,
                        'final_sql': final_sql
                    })
            
            elif parsed['type'] == 'RAW_SQL':
                results.append({
                    'timestamp': timestamp,
                    'client': "{}:{}".format(client_ip, client_port),
                    'type': 'DIRECT',
                    'sql': parsed['query']
                })
                
        except Exception as e:
            print("包解析错误 (Packet #{}): {}".format(pkt.number, e))
    
    return results

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("用法: python3.5 gaussdb_param_parser.py <pcap文件> [端口]")
        sys.exit(1)
    
    pcap_file = sys.argv[1]
    db_port = int(sys.argv[2]) if len(sys.argv) > 2 else 5432
    
    print("正在解析 {} (端口: {})...".format(pcap_file, db_port))
    results = extract_gaussdb_packets(pcap_file, db_port)
    
    # 输出示例
    print("\n[参数化查询示例]")
    param_queries = [r for r in results if r['type'] == 'PARAMETRIC']
    for item in param_queries[:3]:
        print("\n[{}] {}".format(item['timestamp'], item['client']))
        print("预处理语句: {}".format(item['prepared_stmt']))
        print("原始查询: {}".format(item['original_query']))
        print("参数值: {}".format(item['parameters']))
        print("最终SQL: {}".format(item['final_sql'][:100] + "..." if len(item['final_sql']) > 100 else item['final_sql']))
    
    # 保存结果
    with open('gaussdb_param_queries.json', 'w') as f:
        json.dump(results, f, indent=2, ensure_ascii=False)
    print("\n完整结果已保存到 gaussdb_param_queries.json")
